
import torch

# from mvdust3r
def calculate_RRA_RTA(c2w_pred1, c2w_pred2, c2w_gt1, c2w_gt2, eps = 1e-15): # [bs, 3, 3], [bs, 3, 3], [bs, 3, 1], [bs, 3, 1]
    """
    Return:
        RRA: [bs,]
        RTA: [bs,]
    """
    # R1, R2, R1_gt, R2_gt, t1, t2, t1_gt, t2_gt
    R1 = c2w_pred1[:, :3, :3]
    R2 = c2w_pred2[:, :3, :3]
    R1_gt = c2w_gt1[:, :3, :3]
    R2_gt = c2w_gt2[:, :3, :3]
    t1 = c2w_pred1[:, :3, 3:]
    t2 = c2w_pred2[:, :3, 3:]
    t1_gt = c2w_gt1[:, :3, 3:]
    t2_gt = c2w_gt2[:, :3, 3:]

    bs = R1.shape[0]
    R_pred = R1 @ R2.transpose(-1, -2)
    R_gt = R1_gt @ R2_gt.transpose(-1, -2)
    R_diff = R_pred @ R_gt.transpose(-1, -2)

    P_pred_diff = c2w_pred1 @ torch.linalg.inv(c2w_pred2)
    P_gt_diff = c2w_gt1 @ torch.linalg.inv(c2w_gt2)
    # print(R_diff)
    trace_R_diff = R_diff[:,0,0] + R_diff[:,1,1] + R_diff[:,2,2]
    # print(trace_R_diff.shape, 'trace')
    theta = torch.acos((trace_R_diff - 1) / 2)
    theta[theta > torch.pi] = theta[theta > torch.pi] - 2 * torch.pi
    theta[theta < -torch.pi] = theta[theta < -torch.pi] + 2 * torch.pi
    theta = theta * 180 / torch.pi

    # t_pred = t1 - t2
    # t_gt = t1_gt - t2_gt
    # t_pred = P_pred_diff[:, :3, 3:]
    # t_gt = P_gt_diff[:, :3, 3:]
    t_pred = P_pred_diff[:, 3, :3]
    t_gt = P_gt_diff[:, 3, :3]

    # cos_t = (t_pred * t_gt).sum((1,2)) / (torch.norm(t_pred, dim = (-1, -2)) + 1e-8) / (torch.norm(t_gt, dim = (-1, -2)) + 1e-8)
    # theta_t = torch.acos(cos_t)

    default_err = 1e6
    t_norm = torch.norm(t_pred, dim=1, keepdim=True)
    t = t_pred / (t_norm + eps)

    t_gt_norm = torch.norm(t_gt, dim=1, keepdim=True)
    t_gt = t_gt / (t_gt_norm + eps)

    loss_t = torch.clamp_min(1.0 - torch.sum(t * t_gt, dim=1) ** 2, eps)
    theta_t = torch.acos(torch.sqrt(1 - loss_t))

    theta_t[torch.isnan(theta_t) | torch.isinf(theta_t)] = default_err



    theta_t[theta_t > torch.pi] = theta_t[theta_t > torch.pi] - 2 * torch.pi
    theta_t[theta_t < -torch.pi] = theta_t[theta_t < -torch.pi] + 2 * torch.pi
    theta_t = theta_t * 180 / torch.pi

    return theta.abs(), theta_t.abs()

def calibrate_camera_pnpransac(pointclouds, img_points, masks, intrinsics):
    """
    Input:
        pointclouds: (bs, N, 3) 
        img_points: (bs, N, 2) 
    Return:
        rotations: (bs, 3, 3) 
        translations: (bs, 3, 1) 
        c2ws: (bs, 4, 4) 
    """
    bs = pointclouds.shape[0]
    
    camera_matrix = intrinsics.cpu().numpy()  # (bs, 3, 3)
    
    dist_coeffs = np.zeros((5, 1))

    rotations = []
    translations = []
    
    for i in range(bs):
        obj_points = pointclouds[i][masks[i]].cpu().numpy()
        img_pts = img_points[i][[masks[i]]].cpu().numpy()

        success, rvec, tvec, inliers = cv2.solvePnPRansac(obj_points, img_pts, camera_matrix[i], dist_coeffs)

        if success:
            rotation_matrix, _ = cv2.Rodrigues(rvec)
            rotations.append(torch.tensor(rotation_matrix, dtype=torch.float32))
            translations.append(torch.tensor(tvec, dtype=torch.float32))
        else:
            rotations.append(torch.eye(3))
            translations.append(torch.ones(3, 1))

    rotations = torch.stack(rotations).to(pointclouds.device)
    translations = torch.stack(translations).to(pointclouds.device)
    w2cs = torch.eye(4).repeat(bs, 1, 1).to(pointclouds.device)
    w2cs[:, :3, :3] = rotations
    w2cs[:, :3, 3:] = translations
    return torch.linalg.inv(w2cs)

def chamfer_distance(pts1, pts2, mask): # [bs, N, 3], [bs, N, 3], [bs, N]
    bs = pts1.shape[0]
    cd = []
    for i in range(bs):
        disAB = knn_points(pts1[i:i+1][mask[i:i+1]][None], pts2[i:i+1][mask[i:i+1]][None])[0].mean()
        disBA = knn_points(pts2[i:i+1][mask[i:i+1]][None], pts1[i:i+1][mask[i:i+1]][None])[0].mean()
        cd.append(disAB + disBA)
    cd = torch.stack(cd, 0)
    return cd

def rotationInvMSE(pts3d_normalized, gts3d_normalized, mask_all):

    R, sigma, t = umeyama_alignment(pts3d_normalized, gts3d_normalized, mask_all)
    pts3d_normalized_rot = (sigma[:,None,None] * (R @ pts3d_normalized.transpose(-1, -2)).transpose(-1, -2)) + t[:, None] # [bs, h*w, 3]
    local_loss = (pts3d_normalized_rot - gts3d_normalized).norm(dim = -1)[mask_all].mean()
# from evo_utils of CUT3R
def eval_metrics(pred_traj, gt_traj=None, seq="", filename="", sample_stride=1):

    if sample_stride > 1:
        pred_traj[0] = pred_traj[0][::sample_stride]
        pred_traj[1] = pred_traj[1][::sample_stride]
        if gt_traj is not None:
            updated_gt_traj = []
            updated_gt_traj.append(gt_traj[0][::sample_stride])
            updated_gt_traj.append(gt_traj[1][::sample_stride])
            gt_traj = updated_gt_traj

    pred_traj = make_traj(pred_traj)

    if gt_traj is not None:
        gt_traj = make_traj(gt_traj)

        if pred_traj.timestamps.shape[0] == gt_traj.timestamps.shape[0]:
            pred_traj.timestamps = gt_traj.timestamps
        else:
            print(pred_traj.timestamps.shape[0], gt_traj.timestamps.shape[0])

        gt_traj, pred_traj = sync.associate_trajectories(gt_traj, pred_traj)

    # ATE
    traj_ref = gt_traj
    traj_est = pred_traj

    ate_result = main_ape.ape(
        traj_ref,
        traj_est,
        est_name="traj",
        pose_relation=PoseRelation.translation_part,
        align=True,
        correct_scale=True,
    )

    ate = ate_result.stats["rmse"]
    # print(ate_result.np_arrays['error_array'])
    # exit()

    # RPE rotation and translation
    delta_list = [1]
    rpe_rots, rpe_transs = [], []
    for delta in delta_list:
        rpe_rots_result = main_rpe.rpe(
            traj_ref,
            traj_est,
            est_name="traj",
            pose_relation=PoseRelation.rotation_angle_deg,
            align=True,
            correct_scale=True,
            delta=delta,
            delta_unit=Unit.frames,
            rel_delta_tol=0.01,
            all_pairs=True,
        )

        rot = rpe_rots_result.stats["rmse"]
        rpe_rots.append(rot)

    for delta in delta_list:
        rpe_transs_result = main_rpe.rpe(
            traj_ref,
            traj_est,
            est_name="traj",
            pose_relation=PoseRelation.translation_part,
            align=True,
            correct_scale=True,
            delta=delta,
            delta_unit=Unit.frames,
            rel_delta_tol=0.01,
            all_pairs=True,
        )

        trans = rpe_transs_result.stats["rmse"]
        rpe_transs.append(trans)

    rpe_trans, rpe_rot = np.mean(rpe_transs), np.mean(rpe_rots)
    with open(filename, "w+") as f:
        f.write(f"Seq: {seq} \n\n")
        f.write(f"{ate_result}")
        f.write(f"{rpe_rots_result}")
        f.write(f"{rpe_transs_result}")

    print(f"Save results to {filename}")
    return ate, rpe_trans, rpe_rot

def make_traj(args) -> PoseTrajectory3D:
    if isinstance(args, tuple) or isinstance(args, list):
        traj, tstamps = args
        return PoseTrajectory3D(
            positions_xyz=traj[:, :3],
            orientations_quat_wxyz=traj[:, 3:],
            timestamps=tstamps,
        )
    assert isinstance(args, PoseTrajectory3D), type(args)
    return deepcopy(args)